--- +++ @@ -1,5 +1,15 @@ import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react'
 import { Handle, NodeProps } from 'reactflow'
+
+const createHandle = (handleType, handleId, position, ariaLabel, isHidden = true) => (
+  <Handle
+    type={handleType}
+    id={handleId}
+    position={position}
+aria-label={ariaLabel}
+    className={isHidden ? cn(HIDDEN_NODE_CONNECTOR, position === 'left' ? '!left-0' : '!right-0') : ''}
+  />
+);
 
 // Dynamically creates Handle components for nodes
 function createHandleComponent(type, id, position, className) {

@@ -18,6 +28,12 @@  * - sourcePosition (Position): The position of source handles for connecting edges.
  *
  * The component also handles undefined or null data by gracefully not rendering.
+ */
+/**
+ * TableNode Component
+ * Renders a visual representation of a database table as a node within a schema flow.
+ * Depending on whether the table is local or foreign, it delegates to `LocalTableNode` or `ForeignTableNode`.
+ * Handles rendering of table columns, including primary, nullable, unique, and identity indicators.
  */
 export type TableNodeData = {
   name: string;

@@ -45,9 +61,7 @@  * Additionally, it handles rendering connectors (`Handle` components) for establishing relationships between nodes.
  */
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
-  if (!data) {
-    return null;
-  }
+if (!data) return <div>Data is not available</div>;
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }

@@ -59,7 +73,7 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-{targetPosition && createHandleComponent('target', data.name, targetPosition, '!left-0')}
+{targetPosition && createHandle('target', data.name, targetPosition, 'Table connection handle')}
     </header>
   </div>
 );

@@ -123,8 +137,8 @@             {column.format}
           </span>
         </div>
-{targetPosition && createHandleComponent('target', column.id, targetPosition, '!left-0')}
-        {sourcePosition && createHandleComponent('source', column.id, sourcePosition, '!right-0')}
+{targetPosition && createHandle('target', column.id, targetPosition, 'Column connection handle')}
+{sourcePosition && createHandle('source', column.id, sourcePosition, 'Column connection handle')}
       </div>
     ))}
   </div>
   

--- +++ @@ -1,5 +1,16 @@ import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react'
 import { Handle, NodeProps } from 'reactflow'
+/**
+ * `TableNode` Component
+ * Renders a visual representation of a table node within the data schema flow. Supports handling both local
+ * and foreign table nodes. For local table nodes, it displays the table name, columns, and various column properties.
+ * For foreign table nodes, it solely displays the table name. Includes interactive handles for node connections.
+ *
+ * Props:
+ * - `data`: Object containing the table name, if it's a foreign table, and the columns with their properties.
+ * - `targetPosition`: Position of the target handle, if applicable.
+ * - `sourcePosition`: Position of the source handle, if applicable.
+ */
 import { cn } from 'ui'
 import { NODE_WIDTH } from './SchemaFlow.constants'
 
@@ -21,7 +32,22 @@ const ITEM_HEIGHT = 'h-[22px]';
 const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
+const generateHandle = (type, id, position, classNames, ariaLabel) => {
+  return (
+    <Handle
+      type={type}
+      id={id}
+      position={position}
+      className={cn(HIDDEN_NODE_CONNECTOR, classNames)}
+      aria-label={ariaLabel}
+    />
+  );
+};
+
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (!data) {
+    return null;
+  }
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }
@@ -33,13 +59,8 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-      {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
+{targetPosition && (
+        generateHandle("target", data.name, targetPosition, '!left-0', `Target handle for ${data.name}`)
       )}
     </header>
   </div>
@@ -104,21 +125,11 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
+{targetPosition && (
+          generateHandle("target", column.id, targetPosition, '!left-0', `Target handle for ${column.name}`)
         )}
-        {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
+{sourcePosition && (
+          generateHandle("source", column.id, sourcePosition, '!right-0', `Source handle for ${column.name}`)
         )}
       </div>
     ))}
   

--- +++ @@ -3,6 +3,21 @@ import { cn } from 'ui'
 import { NODE_WIDTH } from './SchemaFlow.constants'
 
+const generateHandle = (type, id, position, connectorStyle, label) => (
+  <Handle
+    type={type}
+    id={id}
+    position={position}
+    className={cn(connectorStyle, type === 'target' ? '!left-0' : '!right-0')}
+    aria-label={label}
+  />
+);
+
+/**
+ * The TableNode component renders the visual representation of a table node within the data schema flow.
+ * It supports rendering different layouts and details for foreign and local table nodes,
+ * including table name, columns, and connectors (handles) for schema flow.
+ */
 export type TableNodeData = {
   name: string;
   isForeign: boolean;
@@ -22,6 +37,9 @@ const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (!data) {
+    return <p>Data is not available</p>;
+  }
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }
@@ -33,14 +51,9 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-      {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
-      )}
+{targetPosition && 
+        generateHandle("target", data.name, targetPosition, HIDDEN_NODE_CONNECTOR, "Target connector")
+      }
     </header>
   </div>
 );
@@ -104,22 +117,12 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
-        )}
-        {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
-        )}
+{targetPosition && 
+  generateHandle("target", column.id, targetPosition, HIDDEN_NODE_CONNECTOR, "Target connector")
+}
+        {sourcePosition && 
+          generateHandle("source", column.id, sourcePosition, HIDDEN_NODE_CONNECTOR, "Source connector")
+        }
       </div>
     ))}
   </div>
   

--- +++ @@ -1,15 +1,21 @@ import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react'
 import { Handle, NodeProps } from 'reactflow'
 import { cn } from 'ui'
+/**
+ * `TableNode` Component
+ * Renders a table node within the data schema flow. It differentiates between foreign and local table nodes,
+ * displaying relevant information such as table names, columns, and interactive connectors for schema adjustments.
+ * Useful for visually representing the structure and relationships of database schemas in a user-friendly interface.
+ */
 import { NODE_WIDTH } from './SchemaFlow.constants'
 
-const generateHandle = (type, id, position, connectorStyle, label) => (
+const generateHandle = ({ type, id, position, connectorStyle, ariaLabel }) => (
   <Handle
     type={type}
     id={id}
     position={position}
-    className={cn(connectorStyle, type === 'target' ? '!left-0' : '!right-0')}
-    aria-label={label}
+    className={cn(HIDDEN_NODE_CONNECTOR, connectorStyle)}
+    aria-label={ariaLabel}
   />
 );
 
@@ -37,6 +43,9 @@ const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (data === null || data === undefined) {
+    return null;
+  }
   if (!data) {
     return <p>Data is not available</p>;
   }
@@ -52,7 +61,7 @@     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
 {targetPosition && 
-        generateHandle("target", data.name, targetPosition, HIDDEN_NODE_CONNECTOR, "Target connector")
+generateHandle({ type: "target", id: data.name, position: targetPosition, connectorStyle: HIDDEN_NODE_CONNECTOR, ariaLabel: "Foreign Table Target Connector" })
       }
     </header>
   </div>
@@ -118,11 +127,11 @@           </span>
         </div>
 {targetPosition && 
-  generateHandle("target", column.id, targetPosition, HIDDEN_NODE_CONNECTOR, "Target connector")
+generateHandle({ type: "target", id: column.id, position: targetPosition, connectorStyle: HIDDEN_NODE_CONNECTOR, ariaLabel: `${column.isPrimary ? "Primary Key Connector" : "Column Target Connector"}` })
 }
-        {sourcePosition && 
-          generateHandle("source", column.id, sourcePosition, HIDDEN_NODE_CONNECTOR, "Source connector")
-        }
+{sourcePosition && 
+generateHandle({ type: "source", id: column.id, position: sourcePosition, connectorStyle: HIDDEN_NODE_CONNECTOR, ariaLabel: "Column Source Connector" })
+}
       </div>
     ))}
   </div>
   

--- +++ @@ -1,7 +1,18 @@-import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react'
+import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react';
 import { Handle, NodeProps } from 'reactflow'
+
 import { cn } from 'ui'
 import { NODE_WIDTH } from './SchemaFlow.constants'
+
+const generateHandle = ({ type, id, position, additionalClasses }) => (
+  <Handle
+    type={type}
+    id={id}
+    position={position}
+className={cn(HIDDEN_NODE_CONNECTOR, additionalClasses)}
+    aria-label={`Connect ${type === 'source' ? 'from' : 'to'} table`}
+  />
+);
 
 export type TableNodeData = {
   name: string;
@@ -22,6 +33,9 @@ const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (!data) {
+    return null;
+  }
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }
@@ -33,13 +47,8 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-      {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
+{targetPosition && (
+generateHandle({ type: "target", id: data.name, position: targetPosition, additionalClasses: '!left-0' })
       )}
     </header>
   </div>
@@ -104,21 +113,11 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
+{targetPosition && (
+generateHandle({ type: "target", id: column.id, position: targetPosition, additionalClasses: '!left-0' })
         )}
         {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
+generateHandle({ type: "source", id: column.id, position: sourcePosition, additionalClasses: '!right-0' })
         )}
       </div>
     ))}
   

--- +++ @@ -21,7 +21,15 @@ const ITEM_HEIGHT = 'h-[22px]';
 const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
+/**
+ * TableNode is a React component that renders visual representations of table nodes
+ * within the data schema flow. It supports rendering for both local and foreign table nodes,
+ * displaying table names, columns, and various properties for local tables. It acts as a wrapper
+ * for the LocalTableNode and ForeignTableNode components, dynamically handling data presentation
+ * based on the node type.
+ */
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (!data) return null;
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }
@@ -34,12 +42,13 @@     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
       {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
+<Handle
+  type="target"
+  id={data.name}
+  position={targetPosition}
+  aria-label={`Handle for ${data.name}`}
+  className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
+/>
       )}
     </header>
   </div>
@@ -104,22 +113,24 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
-        )}
-        {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
-        )}
+{targetPosition && (
+<Handle
+  type="target"
+  id={column.id}
+  position={targetPosition}
+  aria-label={`Handle for column ${column.id}`}
+  className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
+/>
+)}
+{sourcePosition && (
+<Handle
+  type="source"
+  id={column.id}
+  position={sourcePosition}
+  aria-label={`Handle for column ${column.id}`}
+  className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
+/>
+)}
       </div>
     ))}
   </div>
   

--- +++ @@ -33,14 +33,7 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-      {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
-      )}
+{targetPosition && renderHandle({ type: "target", id: data.name, position: targetPosition, label: data.name })}
     </header>
   </div>
 );
@@ -104,25 +97,22 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
-        )}
-        {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
-        )}
+{targetPosition && renderHandle({ type: "target", id: column.id, position: targetPosition, label: column.name })}
+{sourcePosition && renderHandle({ type: "source", id: column.id, position: sourcePosition, label: column.name })}
       </div>
     ))}
   </div>
 );
 
+const renderHandle = ({ type, id, position, label }) => {
+  return (
+renderHandle({
+  type: type,
+  id: id,
+  position: position,
+  label: label,
+})
+  );
+};
+
 export default TableNode;
   

@@ -1,16 +1,538 @@
+/**%0A * TableNode Component%0A * This component is responsible for rendering the visual representation of both local and foreign table nodes in the data schema.%0A * It uses the %60isForeign%60 flag in the %60data%60 prop to determine the type of table node and renders accordingly.%0A * For local table nodes, it displays the table name along with columns and their properties.%0A * For foreign table nodes, only the table name is displayed.%0A * It also handles cases where %60data%60 might be null or undefined to prevent runtime errors.%0A */%0A
 import %7B Diamond
@@ -636,16 +636,55 @@
 ctflow'%0A
+import generateHandle from 'reactflow'%0A
 import %7B
@@ -1281,16 +1281,52 @@
 %3E) =%3E %7B%0A
+  if (!data) %7B%0A    return null;%0A  %7D%0A
   if (da
@@ -1768,47 +1768,41 @@
 & (%0A
-        %3C
+generate
 Handle
+(%7B
 %0A
-        
   type
-=
+: 
 %22target%22
 %0A   
@@ -1801,42 +1801,27 @@
 get%22
+,
 %0A  
-        id=%7Bdata.name%7D%0A        
+id: data.name,%0A
   po
@@ -1818,34 +1818,34 @@
 name,%0A  position
-=%7B
+: 
 targetPosition%7D%0A
@@ -1846,79 +1846,54 @@
 tion
-%7D%0A          className=%7Bcn(HIDDEN_NODE_CONNECTOR, '!left-0')%7D%0A        /%3E
+,%0A  label: %60Target handle for $%7Bdata.name%7D%60%0A%7D)
 %0A   
@@ -3966,51 +3966,41 @@
 & (%0A
-          %3C
+generate
 Handle
+(%7B
 %0A
-          
   type
-=
+: 
 %22target%22
 %0A   
@@ -3999,46 +3999,27 @@
 get%22
+,
 %0A  
-          id=%7Bcolumn.id%7D%0A          
+id: column.id,%0A
   po
@@ -4020,26 +4020,26 @@
 ,%0A  position
-=%7B
+: 
 targetPositi
@@ -4044,83 +4044,56 @@
 tion
-%7D%0A            className=%7Bcn(HIDDEN_NODE_CONNECTOR, '!left-0')%7D%0A          /%3E
+,%0A  label: %60Target handle for $%7Bcolumn.name%7D%60%0A%7D)
 %0A   
@@ -4133,51 +4133,41 @@
 & (%0A
-          %3C
+generate
 Handle
+(%7B
 %0A
-          
   type
-=
+: 
 %22source%22
 %0A   
@@ -4166,46 +4166,27 @@
 rce%22
+,
 %0A  
-          id=%7Bcolumn.id%7D%0A          
+id: column.id,%0A
   po
@@ -4187,26 +4187,26 @@
 ,%0A  position
-=%7B
+: 
 sourcePositi
@@ -4211,84 +4211,56 @@
 tion
-%7D%0A            className=%7Bcn(HIDDEN_NODE_CONNECTOR, '!right-0')%7D%0A          /%3E
+,%0A  label: %60Source handle for $%7Bcolumn.name%7D%60%0A%7D)
 %0A   
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -1,3 +1,9 @@+/**
+ * `TableNode` Component
+ * Represents a node in a data schema visualization. This component can render both local and foreign tables.
+ * It displays the table name and, for local tables, a list of columns with icons indicating their properties.
+ * For foreign tables, only the table name is displayed.
+ */
 import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react'
 import { Handle, NodeProps } from 'reactflow'
 import { cn } from 'ui'
@@ -21,7 +27,20 @@ const ITEM_HEIGHT = 'h-[22px]';
 const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
+const renderHandle = (type, id, position, connectorStyle) => (
+  <Handle
+    type={type}
+    id={id}
+    position={position}
+    aria-label={`${type} handle for ${id}`}
+    className={cn(connectorStyle, type === 'target' ? '!left-0' : '!right-0')}
+  />
+);
+
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (!data) {
+    return null;
+  }
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }
@@ -33,14 +52,7 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-      {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
-      )}
+{targetPosition && renderHandle('target', data.name, targetPosition, HIDDEN_NODE_CONNECTOR)}
     </header>
   </div>
 );
@@ -104,22 +116,8 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
-        )}
-        {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
-        )}
+{targetPosition && renderHandle('target', column.id, targetPosition, HIDDEN_NODE_CONNECTOR)}
+        {sourcePosition && renderHandle('source', column.id, sourcePosition, HIDDEN_NODE_CONNECTOR)}
       </div>
     ))}
   </div>
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -21,7 +21,15 @@ const ITEM_HEIGHT = 'h-[22px]';
 const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
+/**
+ * `TableNode` Component
+ * Represents a node within the data schema visualization, handling both local and foreign table types.
+ * It renders the table name and, for local tables, the columns with icons indicating their properties.
+ * For foreign tables, only the table name is rendered. Also facilitates connections between nodes in the schema flow.
+ */
+
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (!data) return null; // This line checks if `data` is null or undefined, and if so, prevents the rest of the component from rendering.
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }
@@ -33,21 +41,18 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-      {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
-      )}
+{targetPosition && (
+  <>{renderHandle({type: "target", id: data.name, position: targetPosition, label: `Target handle for ${data.name}`})}</>
+)}
     </header>
   </div>
 );
 
 
 
-const LocalTableNode = ({ data, targetPosition, sourcePosition }) => (
+const LocalTableNode = ({ data, targetPosition, sourcePosition }) => {
+  if (!data) return null;
+  return (
   <div
     className="border border-[0.5px] overflow-hidden rounded-[4px] shadow-sm"
     style={{ width: NODE_WIDTH / 2 }}
@@ -104,25 +109,27 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
-        )}
-        {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
-        )}
+{targetPosition && (
+  <>{renderHandle({type: "target", id: column.id, position: targetPosition, label: `Target handle for column ${column.name}`})}</>
+)}
+{sourcePosition && (
+  <>{renderHandle({type: "source", id: column.id, position: sourcePosition, label: `Source handle for column ${column.name}`})}</>
+)}
       </div>
     ))}
   </div>
 );
 
+}
+
+const renderHandle = ({type, id, position, label}) => (
+  <Handle
+    type={type}
+    id={id}
+    position={position}
+    aria-label={label}
+    className={cn(HIDDEN_NODE_CONNECTOR, position === 'left' ? '!left-0' : '!right-0')}
+  />
+);
+
 export default TableNode;
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -1,6 +1,12 @@ import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react'
 import { Handle, NodeProps } from 'reactflow'
 import { cn } from 'ui'
+/**
+ * The TableNode component renders a visual representation of a table node within a data schema flow.
+ * It supports rendering both local and foreign table nodes with distinctions based on the 'isForeign' property from its 'data' prop.
+ * For foreign table nodes, only the table name is displayed, while local nodes show detailed column information.
+ * This component also safely handles null or undefined 'data' props to prevent runtime errors.
+ */
 import { NODE_WIDTH } from './SchemaFlow.constants'
 
 export type TableNodeData = {
@@ -22,6 +28,9 @@ const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (!data) {
+    return null; // Or any fallback UI component
+  }
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }
@@ -33,14 +42,7 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-      {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
-      )}
+{targetPosition && <Handle type="target" id={data.name} position={targetPosition} aria-label={`Connect to ${data.name}`} />}
     </header>
   </div>
 );
@@ -104,22 +106,8 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
-        )}
-        {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
-        )}
+{targetPosition && <Handle type="target" id={column.id} position={targetPosition} aria-label={`Connect to column ${column.id}`} />}
+{sourcePosition && <Handle type="source" id={column.id} position={sourcePosition} aria-label={`Connect from column ${column.id}`} />}
       </div>
     ))}
   </div>
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -1,7 +1,26 @@+/**
+ * `TableNode` component represents a visual node within a schema flow diagram.
+ * It can render both local and foreign table nodes, displaying table names, columns,
+ * and various column properties using icons. The component supports interactive elements
+ * for establishing connections between nodes within the schema flow.
+ * 
+ * It can gracefully handle `null` or `undefined` `data`, ensuring the component's robustness
+ * in various usage scenarios.
+ */
 import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react'
 import { Handle, NodeProps } from 'reactflow'
 import { cn } from 'ui'
 import { NODE_WIDTH } from './SchemaFlow.constants'
+
+const generateHandle = (id: string, position: 'target' | 'source', positionClass: string) => (
+  <Handle
+    type={position}
+    id={id}
+    position={position}
+    className={positionClass}
+    aria-label={`${position} handle for ${id}`} // Enhance accessibility
+  />
+);
 
 export type TableNodeData = {
   name: string;
@@ -22,6 +41,9 @@ const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
 const TableNode = ({ data, targetPosition, sourcePosition }: NodeProps<TableNodeData>) => {
+  if (!data) {
+    return null;  // Return `null` if `data` is null or undefined.
+  }
   if (data.isForeign) {
     return <ForeignTableNode data={data} targetPosition={targetPosition} />;
   }
@@ -33,14 +55,7 @@   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
       {data.name}
-      {targetPosition && (
-        <Handle
-          type="target"
-          id={data.name}
-          position={targetPosition}
-          className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-        />
-      )}
+{targetPosition && generateHandle(data.name, 'target', cn(HIDDEN_NODE_CONNECTOR, '!left-0'))}
     </header>
   </div>
 );
@@ -104,22 +119,8 @@             {column.format}
           </span>
         </div>
-        {targetPosition && (
-          <Handle
-            type="target"
-            id={column.id}
-            position={targetPosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!left-0')}
-          />
-        )}
-        {sourcePosition && (
-          <Handle
-            type="source"
-            id={column.id}
-            position={sourcePosition}
-            className={cn(HIDDEN_NODE_CONNECTOR, '!right-0')}
-          />
-        )}
+{targetPosition && generateHandle(column.id, 'target', cn(HIDDEN_NODE_CONNECTOR, '!left-0'))}
+{sourcePosition && generateHandle(column.id, 'source', cn(HIDDEN_NODE_CONNECTOR, '!right-0'))}
       </div>
     ))}
   </div>
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -2,6 +2,11 @@ import { Handle, NodeProps } from 'reactflow'
 import { cn } from 'ui'
 import { NODE_WIDTH } from './SchemaFlow.constants'
+
+// Constants for node styles
+const HIDDEN_NODE_CONNECTOR = '!h-px !w-px !min-w-0 !min-h-0 !cursor-grab !border-0 !opacity-0';
+const ITEM_HEIGHT = 'h-[22px]';
+const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
 const generateHandle = (id: string, position: 'target' | 'source', positionClass: string) => (
   <Handle
@@ -26,10 +31,6 @@     format: string;
   }>;
 };
-
-const HIDDEN_NODE_CONNECTOR = '!h-px !w-px !min-w-0 !min-h-0 !cursor-grab !border-0 !opacity-0';
-const ITEM_HEIGHT = 'h-[22px]';
-const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
 /**
  * `TableNode` component represents a visual node within a schema flow diagram.
@@ -62,6 +63,19 @@ 
 
 
+/**
+ * Renders a local table node within a schema flow diagram.
+ *
+ * This component is responsible for visually representing a local table, including its name, columns,
+ * and their respective properties. It supports interactive elements such as handles, allowing connections
+ * between nodes within the schema flow based on the table's relationships.
+ *
+ * @param {TableNodeData} data - An object containing the table's name, an indication if it is a foreign table,
+ * and an array of column objects detailing the column name, properties and format.
+ * @param {'top' | 'bottom' | 'left' | 'right' | undefined} targetPosition - Designates the position of the target handle within the node for incoming connections.
+ * @param {'top' | 'bottom' | 'left' | 'right' | undefined} sourcePosition - Specifies the position of the source handle within the node for outgoing connections.
+ * @returns A JSX element representing the local table node with appropriately placed handles for establishing connections, if `targetPosition` or `sourcePosition` are provided.
+ */
 const LocalTableNode = ({ data, targetPosition, sourcePosition }) => (
   <div
     className="border border-[0.5px] overflow-hidden rounded-[4px] shadow-sm"
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -8,6 +8,16 @@ const ITEM_HEIGHT = 'h-[22px]';
 const FLEX_SHRINK_TEXT_LIGHT = 'flex-shrink-0 text-light';
 
+/**
+ * Creates a `Handle` component for `TableNode` used for establishing node connections in a schema flow diagram.
+ * This handle is visualized based on its `position` within the parent node and customized with specific classes for styling.
+ * 
+ * @param {string} id - A unique identifier for the handle, typically related to the column or table it represents.
+ * @param {'target' | 'source'} position - Specifies the handle type, determining whether it's for incoming (`target`) or outgoing (`source`) connections.
+ * @param {string} positionClass - Additional CSS classes to customize the appearance and positioning of the handle in the UI.
+ * 
+ * @returns A `Handle` component configured with an `aria-label` for accessibility and designated styling for integration into the schema flow diagram.
+ */
 const generateHandle = (id: string, position: 'target' | 'source', positionClass: string) => (
   <Handle
     type={position}
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -62,6 +62,16 @@   return <LocalTableNode data={data} targetPosition={targetPosition} sourcePosition={sourcePosition} />;
 };
 
+/**
+ * Renders a foreign table node within a schema flow diagram.
+ *
+ * This component is responsible for visually representing a foreign table, displaying its name.
+ * It supports an interactive element such as a handle, allowing connections between nodes
+ * within the schema flow based on the table's relationships.
+ *
+ * @param {TableNodeData} data - An object containing the table's name and indicating if it is a foreign table.
+ * @param {'top' | 'bottom' | 'left' | 'right' | undefined} targetPosition - Designates the position of the target handle within the node for incoming connections.
+ */
 const ForeignTableNode = ({ data, targetPosition }) => (
   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -28,6 +28,17 @@   />
 );
 
+/**
+ * Renders a foreign table node within a schema flow diagram.
+ *
+ * This component is responsible for visually representing a foreign table node, displaying the table name
+ * and incorporating a specialized handle for establishing target connections within the schema flow based on
+ * the table's foreign key relationships.
+ *
+ * @param {TableNodeData} data - An object containing the name of the foreign table and a list of its columns, detailing the column name, properties, and format.
+ * @param {'top' | 'bottom' | 'left' | 'right' | undefined} targetPosition - Designates the position of the target handle within the node for incoming connections.
+ * @returns A JSX element representing the foreign table node with an appropriately placed handle for establishing connections if `targetPosition` is provided.
+ */
 export type TableNodeData = {
   name: string;
   isForeign: boolean;
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -62,6 +62,19 @@   return <LocalTableNode data={data} targetPosition={targetPosition} sourcePosition={sourcePosition} />;
 };
 
+/**
+ * Renders a foreign table node within a schema flow diagram.
+ *
+ * This component visualizes a foreign table node, showcasing its name and enabling
+ * the establishment of incoming connections through a visually represented handle.
+ * The appearance and position of the handle are adjusted based on the `targetPosition`
+ * provided, facilitating the intuitive mapping of relationships in the schema flow.
+ *
+ * @param {Object} props - The properties passed to the component.
+ * @param {TableNodeData} props.data - An object containing the table's name and column details.
+ * @param {'top' | 'bottom' | 'left' | 'right'} props.targetPosition - Specifies the position of the target handle for incoming connections.
+ * @returns A JSX element representing the foreign table node, including a positional handle for establishing connections within the schema flow diagram.
+ */
 const ForeignTableNode = ({ data, targetPosition }) => (
   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -62,6 +62,18 @@   return <LocalTableNode data={data} targetPosition={targetPosition} sourcePosition={sourcePosition} />;
 };
 
+/**
+ * Renders a foreign table node within a schema flow diagram.
+ *
+ * This component is responsible for visually representing a foreign table, displaying the table's name
+ * and allowing the connection from this node to others within the schema flow through the use of handles.
+ * It is visually differentiated and simplified compared to local table nodes to emphasize its external nature.
+ *
+ * @param {TableNodeData} data - An object containing the table's name, indicating it as a foreign table,
+ * and potentially including an array of column objects (not used in this component).
+ * @param {'top' | 'bottom' | 'left' | 'right' | undefined} targetPosition - Designates the position of the target handle within the node for incoming connections.
+ * @returns A JSX element representing the foreign table node with an appropriately placed handle for incoming connections, if `targetPosition` is provided.
+ */
 const ForeignTableNode = ({ data, targetPosition }) => (
   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">
   

--- codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx+++ codebase/apps/database-new/app/[threadId]/[runId]/[messageId]/TableNode.tsx@@ -62,6 +62,17 @@   return <LocalTableNode data={data} targetPosition={targetPosition} sourcePosition={sourcePosition} />;
 };
 
+/**
+ * Renders a foreign table node within a schema flow diagram.
+ *
+ * This component is responsible for visually representing a foreign table, indicating its name and allowing
+ * connections to other nodes within the schema flow by leveraging the `generateHandle` function based on the table's
+ * relationships. This functionality facilitates the establishment of inbound connections to the foreign table node.
+ *
+ * @param {TableNodeData} data - An object containing the table's name and various properties, including an indication if it is a foreign table, and an array of column objects detailing the column names, properties, and formats.
+ * @param {'top' | 'bottom' | 'left' | 'right' | undefined} targetPosition - Designates the position of the target handle within the node for incoming connections.
+ * @returns A JSX element representing the foreign table node with an appropriately placed handle for establishing connections, if `targetPosition` is provided.
+ */
 const ForeignTableNode = ({ data, targetPosition }) => (
   <div className="rounded-lg">
     <header className="text-[0.5rem] leading-5 font-bold px-2 text-center bg-brand text-gray-300">